import math
import time
import sys
import os
from getDCSdata import get_dcs_data
import keyboard_control as kc

class RunwayDecelerationController:
    def __init__(self):
        """初始化跑道减速控制器"""
        self.target_heading = None
        self.last_operation_time = time.time()
        self.operation_interval = 0.5  # 操作间隔，单位秒
        self.max_retry_count = 3  # 最大重试次数
        self.retry_delay = 2  # 重试延迟，单位秒
        self.last_ias = 0  # 上一时刻的空气速度
        self.last_time = time.time()  # 上一时刻的时间
        self.acceleration_threshold = 1.5  # 加速度阈值 (m/s²)，可以调整
        self.last_status = None  # 用于保存上一时刻状态

    def activate_dcs_window(self):
        """激活 DCS 窗口"""
        for try_count in range(self.max_retry_count):
            try:
                kc.activate_dcs_window()
                print(f"DCS 窗口激活成功 (尝试 {try_count + 1}/{self.max_retry_count})")
                return True
            except Exception as error:
                print(f"DCS 窗口激活失败 (尝试 {try_count + 1}/{self.max_retry_count}): {str(error)}")
                if try_count < self.max_retry_count - 1:
                    time.sleep(self.retry_delay)
        return False

    def get_current_status(self):
        """获取并验证飞行状态数据"""
        try:
            status = get_dcs_data()
            if status is None:
                print("收到空的状态数据")
                return None

            # 设置跑道方向（以当前航向作为初始跑道方向）
            if self.target_heading is None:
                self.target_heading = status['heading_deg']
                print(f"初始目标航向设置为: {self.target_heading:.1f}°")
            status['RunwayHeading'] = self.target_heading

            # 验证必要字段
            necessary_fields = ['IAS', 'LLA_balt', 'heading_deg', 'Gears', 'Flaps', 'AirBrake', 'VV_0']
            for field in necessary_fields:
                if field not in status:
                    print(f"状态数据缺少必要字段: {field}")
                    return None

            # 调试输出
            print(
                f"状态更新 | 速度: {status['IAS']:.1f} m/s | 高度: {status['LLA_balt']:.1f} m | 航向: {status['heading_deg']:.1f}°")
            return status

        except Exception as error:
            print(f"获取状态数据时发生异常: {str(error)}")
            return None

    def check_runway_status(self, current_status):
        """增强的跑道状态检查逻辑"""
        # 添加垂直速度检查
        if abs(current_status['VV_0']) > 20:  # 垂直速度过大
            print(f"垂直速度异常: {current_status['VV_0']:.1f} m/s")
            return False

        # 高度检查（单位：米）
        if current_status['LLA_balt'] > 500:
            print("跑道检查: 高度超过 500 米，不执行控制")
            return False

        # 速度检查（单位：米每秒）
        if current_status['IAS'] > 200:
            print("跑道检查: 速度超过 100 米/秒，不执行控制")
            return False

        # 起落架检查（范围 0-1）
        if current_status.get('Gears', 0) < 0.8:
            print(f"跑道检查: 起落架未部署 ({current_status['Gears']:.2f})")
            return False

        # 襟翼检查（范围 0-1）
        if current_status.get('Flaps', 0) < 0.5:
            print(f"跑道检查: 襟翼未放下 ({current_status['Flaps']:.2f})")
            return False

        print("跑道检查: 符合跑道状态条件")
        return True

    def calculate_acceleration(self, current_ias):
        """带平滑处理的加速度计算"""
        current_time = time.time()
        time_diff = current_time - self.last_time

        if time_diff < 0.1:  # 最小时间间隔
            return getattr(self, 'smoothed_accel', 0)

        # 原始加速度计算
        raw_accel = (current_ias - self.last_ias) / time_diff

        # 指数移动平均平滑
        self.smoothed_accel = 0.6 * getattr(self, 'smoothed_accel', 0) + 0.4 * raw_accel

        # 更新状态
        self.last_ias = current_ias
        self.last_time = current_time

        return self.smoothed_accel

    def perform_deceleration_control(self, current_status):
        """执行减速和方向控制"""
        current_time = time.time()
        if current_time - self.last_operation_time < self.operation_interval:
            return

        # 保存前一时刻状态用于奖励计算
        prev_status = getattr(self, 'last_status', None)
        self.last_status = current_status.copy()

        airspeed = current_status["IAS"]
        heading_error = abs((current_status["heading_deg"] - current_status["RunwayHeading"] + 180) % 360 - 180)

        # 计算加速度
        acceleration = self.calculate_acceleration(airspeed)

        # 计算并打印奖励
        if prev_status is not None:
            reward, reward_info = calculate_runway_deceleration_reward(prev_status, current_status, None)
            print("\n=== 奖励信息 ===")
            print(f"总奖励: {reward:.2f}")
            for key, value in reward_info.items():
                print(f"{key}: {value:.2f}")
            print("===============\n")

        # 判断加速度是否超过阈值，执行减速操作
        if abs(acceleration) > self.acceleration_threshold:
            print(f"加速度超过阈值 {self.acceleration_threshold} m/s², 当前加速度: {acceleration:.2f} m/s²")
            self.apply_brake_for_acceleration()

        # 航向修正
        if heading_error > 5:
            if (current_status["heading_deg"] - current_status["RunwayHeading"] + 180) % 360 - 180 > 0:
                kc.foot_left(duration=0.2)
                print("航向修正: 向左转")
            else:
                kc.foot_right(duration=0.2)
                print("航向修正: 向右转")

        # 处理爬升阶段的减速控制
        if current_status.get('VV_0', 0) > 1.0:  # 判断是否在爬升阶段
            print(f"正在爬升阶段，垂直速度: {current_status['VV_0']:.2f} m/s")
            if airspeed > 60:
                if current_status["AirBrake"] < 0.5:
                    kc.airbrake_control(duration=0.3)
                    print("爬升阶段减速: 部分放出空中刹车")
                kc.wheel_control(duration=0.3)
                print("爬升阶段减速: 轻微刹车")
            elif airspeed > 30:
                if current_status["AirBrake"] < 0.9:
                    kc.airbrake_control(duration=0.5)
                    print("爬升阶段减速: 完全放出空中刹车")
                kc.wheel_control(duration=0.5)
                print("爬升阶段减速: 中等刹车")
            else:
                kc.wheel_control(duration=1.0)
                print("爬升阶段减速: 全力刹车")
        else:
            # 如果不在爬升阶段，进行正常的减速操作
            if airspeed > 60:
                if current_status["AirBrake"] < 0.5:
                    kc.airbrake_control(duration=0.3)
                    print("减速操作: 部分放出空中刹车")
                kc.wheel_control(duration=0.3)
                print("减速操作: 轻微刹车")
            elif airspeed > 30:
                if current_status["AirBrake"] < 0.9:
                    kc.airbrake_control(duration=0.5)
                    print("减速操作: 完全放出空中刹车")
                kc.wheel_control(duration=0.5)
                print("减速操作: 中等刹车")
            else:
                kc.wheel_control(duration=1.0)
                print("减速操作: 全力刹车")

        self.last_operation_time = current_time

    def perform_deceleration_control(self, current_status):
        """执行完整的减速控制流程"""
        current_time = time.time()
        if current_time - self.last_operation_time < self.operation_interval:
            return

        # 保存前一时刻状态用于奖励计算
        prev_status = getattr(self, 'last_status', None)
        self.last_status = current_status.copy()

        # 获取关键状态参数
        airspeed = current_status["IAS"]
        altitude = current_status["LLA_balt"]
        vertical_speed = current_status["VV_0"]
        heading_error = abs((current_status["heading_deg"] - current_status["RunwayHeading"] + 180) % 360 - 180)
        lateral_offset = abs(current_status.get("RunwayOffset", 0))

        # 计算加速度（带平滑处理）
        acceleration = self.calculate_acceleration(airspeed)

        # ===== 1. 动态刹车控制策略 =====
        brake_intensity = 0  # 刹车强度 0-1
        airbrake_action = False  # 是否使用空气刹车

        # 分速度阶段控制
        if airspeed > 30:  # 高速阶段
            if acceleration > 1.0:  # 急减速需求
                brake_intensity = min(0.8, 0.3 + acceleration / 5)
                airbrake_action = current_status["AirBrake"] < 0.9
            else:
                brake_intensity = 0.4
                airbrake_action = current_status["AirBrake"] < 0.7

        elif airspeed > 15:  # 中速阶段
            brake_intensity = 0.6
            airbrake_action = current_status["AirBrake"] < 0.5

        else:  # 低速阶段
            if airspeed > 5:  # 低速滑行
                brake_intensity = 0.3
            else:  # 即将停止
                brake_intensity = 1.0

        # 执行控制动作
        if airbrake_action:
            kc.airbrake_control(duration=0.5)
            print(f"空气刹车控制 | 当前状态: {current_status['AirBrake']:.1f}")

        if brake_intensity > 0:
            kc.wheel_control(duration=brake_intensity)
            print(f"机轮刹车控制 | 强度: {brake_intensity:.1f} | 速度: {airspeed:.1f}m/s")

        # ===== 2. 航向修正 =====
        if heading_error > 5:  # 需要航向修正
            correction_duration = min(0.5, heading_error / 30)  # 动态修正时长

            if (current_status["heading_deg"] - current_status["RunwayHeading"] + 180) % 360 - 180 > 0:
                kc.foot_left(duration=correction_duration)
                print(f"航向修正: 左转 {correction_duration:.1f}s | 偏差: {heading_error:.1f}°")
            else:
                kc.foot_right(duration=correction_duration)
                print(f"航向修正: 右转 {correction_duration:.1f}s | 偏差: {heading_error:.1f}°")

        # ===== 3. 奖励计算与反馈 =====
        if prev_status is not None:
            reward, reward_info = calculate_runway_deceleration_reward(prev_status, current_status, None)

            print("\n=== 性能反馈 ===")
            print(f"速度变化: {reward_info['speed_change']:.2f} m/s | 减速强度: {reward_info['decel_intensity']:.2f}")
            print(f"刹车奖励: {reward_info['brake_reward']:.2f} | 航向奖励: {reward_info['alignment_reward']:.2f}")
            print(f"总奖励: {reward:.2f}/1.0")
            print("===============")

        # ===== 4. 状态检查 =====
        if airspeed < 1 and altitude < 5:
            print("飞机已完全停稳！")
            return True  # 返回True表示任务完成

        self.last_operation_time = current_time
        return False  # 返回False表示需要继续控制

    def apply_brake_for_acceleration(self):
        """根据加速度过大时应用刹车"""
        print("检测到加速度过大，执行减速操作...")
        # 放出空中刹车并轻微刹车
        kc.airbrake_control(duration=0.5)
        kc.wheel_control(duration=0.5)
        print("加速过快，已采取减速操作")

    def run(self):
        """主控制循环"""
        print("==== 跑道减速控制系统启动 ====")

        if not self.activate_dcs_window():
            return False

        try:
            while True:
                current_status = self.get_current_status()

                if current_status is None:
                    print("等待有效的状态数据...")
                    time.sleep(1)
                    continue

                # 判断是否满足控制条件：高度 < 500 米 且 速度 < 200 米/秒
                if self.check_runway_status(current_status):
                    if current_status['IAS'] > 5:
                        self.perform_deceleration_control(current_status)
                    else:
                        print("飞机已完全停稳")
                        break
                else:
                    print("飞机不处于跑道减速状态")

                time.sleep(0.1)

        except KeyboardInterrupt:
            print("\n用户中断控制")
        except Exception as error:
            print(f"控制循环发生异常: {str(error)}")
            return False
        finally:
            print("==== 系统关闭 ====")
        return True
def calculate_runway_deceleration_reward(prev_state, current_state, action):

    # 状态参数提取
    ias = current_state["IAS"]  # 当前空速(m/s)
    prev_ias = prev_state["IAS"] if prev_state else ias
    alt = current_state["LLA_balt"]  # 高度(m)
    wheel_brake = current_state.get("WheelBrake", 0)  # 刹车状态(0-1)
    airbrake = current_state["AirBrake"]  # 扰流板状态(0-1)
    heading_error = abs((current_state["heading_deg"] - current_state["RunwayHeading"] + 180) % 360 - 180)
    lateral_offset = abs(current_state.get("RunwayOffset", 0))  # 跑道横向偏移(m)

    # 物理参数(示例值，需根据实际机型调整)
    MAX_BRAKE_FORCE = 1.0  # 最大刹车力(标准化)
    BRAKE_EFFICIENCY = 0.8  # 刹车效率系数
    AERO_DRAG_COEFF = 0.2  # 气动阻力系数

    # 1. 核心减速效率奖励（权重0.4）
    # 计算理论最大可能减速度
    max_decel = (MAX_BRAKE_FORCE * BRAKE_EFFICIENCY +
                 airbrake * AERO_DRAG_COEFF * ias ** 2 / 100)

    # 实际减速度
    actual_decel = (prev_ias - ias)  # 正值表示减速

    # 减速效率(实际减速度/理论最大减速度)
    decel_efficiency = min(actual_decel / max_decel, 1.0) if max_decel > 0 else 0

    # 分段奖励策略
    if ias > 60:  # 高速阶段(60m/s以上)
        decel_reward = 0.4 * 0.3 * decel_efficiency  # 权重0.4中的30%
    elif ias > 30:  # 中速阶段(30-60m/s)
        decel_reward = 0.4 * 0.6 * decel_efficiency  # 权重0.4中的60%
    else:  # 低速阶段(30m/s以下)
        decel_reward = 0.4 * 1.0 * decel_efficiency  # 权重0.4中的100%

    # 2. 能量管理奖励（权重0.2）
    # 扰流板使用奖励
    spoiler_reward = 0.2 * 0.5 * airbrake if ias > 30 else 0  # 权重0.2中的50%
    # 刹车使用奖励
    brake_reward = 0.2 * 0.5 * wheel_brake if ias < 60 else 0  # 权重0.2中的50%

    # 3. 跑道保持奖励（权重0.3）
    # 航向对准奖励(指数衰减)
    alignment_reward = 0.3 * 0.6 * (math.exp(-heading_error / 5.0) - 0.5)  # 权重0.3中的60%
    # 横向偏移惩罚(二次增长)
    offset_penalty = 0.3 * 0.4 * (-0.02 * lateral_offset ** 2)  # 权重0.3中的40%

    # 4. 稳定性奖励（权重0.1）
    pitch = abs(current_state["pitch_deg"])
    bank = abs(current_state["bank_deg"])
    stability_reward = 0.1 * (-0.05 * (pitch + bank))  # 权重0.1

    # 5. 特殊状态奖励/惩罚（额外加减分，但不超过总奖励范围）
    # 完全停止奖励（最高0.2）
    stop_bonus = min(0.2, 0.2 * (1.0 if ias < 1 and alt < 1 and lateral_offset < 5 else 0))

    # 跑道偏离惩罚（最多-0.2）
    runway_penalty = max(-0.2, -0.2 if lateral_offset > 15 else 0)

    # 刹车锁死惩罚（最多-0.1）
    if (wheel_brake > 0.9 and ias > 20 and (prev_ias - ias) < 0.5):
        skid_penalty = max(-0.1, -0.1)
    else:
        skid_penalty = 0

    # 综合奖励计算（确保在0-1范围内）
    total_reward = max(0, min(1, (
            decel_reward +
            spoiler_reward +
            brake_reward +
            alignment_reward +
            offset_penalty +
            stability_reward +
            stop_bonus +
            runway_penalty +
            skid_penalty
    )))

    # 调试信息
    info = {
        "decel_reward": decel_reward,
        "spoiler_reward": spoiler_reward,
        "brake_reward": brake_reward,
        "alignment_reward": alignment_reward,
        "offset_penalty": offset_penalty,
        "stability_reward": stability_reward,
        "stop_bonus": stop_bonus,
        "runway_penalty": runway_penalty,
        "skid_penalty": skid_penalty,
        "decel_efficiency": decel_efficiency
    }

    return total_reward, info

if __name__ == "__main__":
    controller = RunwayDecelerationController()

    try:
        success = controller.run()
        if not success:
            sys.exit(1)
    except Exception as error:
        print(f"发生未捕获的异常: {str(error)}")
        sys.exit(1)
